#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

f = xr.open_dataset('swm_mz.1440x720.h0.nc')

fig = plt.figure(figsize=(10, 8))
proj = ccrs.PlateCarree()

ax = fig.add_subplot(2, 2, 1, projection=proj)
plot_contour_map(ax, f.u[args.time_step,0,:,:], cmap='jet', add_cyclic_point=False)

ax = fig.add_subplot(2, 2, 2, projection=proj)
plot_contour_map(ax, f.v[args.time_step,0,:,:], cmap='jet', add_cyclic_point=False)

ax = fig.add_subplot(2, 2, 3, projection=proj)
plot_contour_map(ax, f.gz[args.time_step,0,:,:], cmap='jet')

ax = fig.add_subplot(2, 2, 4, projection=proj)
plot_contour_map(ax, f.pv[args.time_step,0,:,:], cmap='jet')

plt.tight_layout()
if args.output:
  plt.savefig(args.output, dpi=300)
else:
  plt.show()
